#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
点云数据可视化脚本
读取PTS格式点云数据和对应的SEG分类文件，并根据分类进行着色显示
"""

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
plt.rcParams['font.sans-serif'] = ['Microsoft YaHei', 'Noto Sans CJK SC']
plt.rcParams['axes.unicode_minus']=False #用来正常显示负号

def read_pts_file(pts_path):
    """
    读取PTS格式点云文件
    
    Args:
        pts_path: PTS文件路径
        
    Returns:
        numpy数组，shape为(N, 3)，每行为[x, y, z]坐标
    """
    points = []
    with open(pts_path, 'r') as f:
        for line in f:
            # 去除空白字符并分割
            coords = line.strip().split()
            if len(coords) == 3:
                # 转换为浮点数
                x, z, y = float(coords[0]), float(coords[1]), float(coords[2])
                points.append([x, y, z])
    
    return np.array(points)


def read_seg_file(seg_path):
    """
    读取SEG格式分类文件
    
    Args:
        seg_path: SEG文件路径
        
    Returns:
        numpy数组，每个元素为对应点的分类标签（整数）
    """
    labels = []
    with open(seg_path, 'r') as f:
        for line in f:
            label = int(line.strip())
            labels.append(label)
    
    return np.array(labels)


def visualize_pointcloud(points, labels):
    """
    可视化点云数据，根据分类标签着色
    
    Args:
        points: 点云坐标数组，shape为(N, 3)
        labels: 分类标签数组，shape为(N,)
    """
    # 创建3D图形
    fig = plt.figure(figsize=(12, 10))
    ax = fig.add_subplot(111, projection='3d')
    
    # 定义颜色映射：1=红色，2=绿色，3=蓝色
    color_map = {
        1: 'red',
        2: 'green',
        3: 'blue'
    }
    
    # 获取唯一的标签值
    unique_labels = np.unique(labels)
    
    # 为每个类别分别绘制点云
    for label in unique_labels:
        # 筛选当前类别的点
        mask = labels == label
        points_subset = points[mask]
        
        # 获取对应颜色
        color = color_map.get(label, 'gray')  # 如果标签不在1-3范围内，使用灰色
        
        # 绘制散点
        ax.scatter(
            points_subset[:, 0],  # X坐标
            points_subset[:, 1],  # Y坐标
            points_subset[:, 2],  # Z坐标
            c=color,
            s=5,  # 点的大小
            alpha=0.6,  # 透明度
            label=f'类别 {label}'
        )
    
    # 设置坐标轴标签
    ax.set_xlabel('X轴', fontsize=12)
    ax.set_ylabel('Y轴', fontsize=12)
    ax.set_zlabel('Z轴', fontsize=12)
    
    # 设置标题
    ax.set_title('点云分类可视化', fontsize=14, fontweight='bold')
    
    # 添加图例
    ax.legend(loc='upper right', fontsize=10)
    
    # 设置坐标轴比例相等，使显示更真实
    max_range = np.array([
        points[:, 0].max() - points[:, 0].min(),
        points[:, 1].max() - points[:, 1].min(),
        points[:, 2].max() - points[:, 2].min()
    ]).max() / 2.0
    
    mid_x = (points[:, 0].max() + points[:, 0].min()) * 0.5
    mid_y = (points[:, 1].max() + points[:, 1].min()) * 0.5
    mid_z = (points[:, 2].max() + points[:, 2].min()) * 0.5
    
    ax.set_xlim(mid_x - max_range, mid_x + max_range)
    ax.set_ylim(mid_y - max_range, mid_y + max_range)
    ax.set_zlim(mid_z - max_range, mid_z + max_range)
    
    # 调整视角
    ax.view_init(elev=20, azim=45)
    
    # 显示网格
    ax.grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.show()


def main():
    """主函数"""
    # 定义文件路径
    pts_file = 'demo_data/pst_seg/02691156/points/1a04e3eab45ca15dd86060f189eb133.pts'
    seg_file = 'demo_data/pst_seg/02691156/points_label/1a04e3eab45ca15dd86060f189eb133.seg'
    
    print("正在读取点云数据...")
    points = read_pts_file(pts_file)
    print(f"成功读取 {len(points)} 个点")
    
    print("正在读取分类标签...")
    labels = read_seg_file(seg_file)
    print(f"成功读取 {len(labels)} 个标签")
    
    # 验证点数和标签数是否匹配
    if len(points) != len(labels):
        print(f"警告：点数({len(points)})与标签数({len(labels)})不匹配！")
        return
    
    # 显示统计信息
    print("\n点云统计信息:")
    print(f"  X范围: [{points[:, 0].min():.6f}, {points[:, 0].max():.6f}]")
    print(f"  Y范围: [{points[:, 1].min():.6f}, {points[:, 1].max():.6f}]")
    print(f"  Z范围: [{points[:, 2].min():.6f}, {points[:, 2].max():.6f}]")
    
    print("\n分类统计:")
    unique_labels, counts = np.unique(labels, return_counts=True)
    for label, count in zip(unique_labels, counts):
        print(f"  类别 {label}: {count} 个点 ({count/len(labels)*100:.2f}%)")
    
    print("\n正在生成可视化...")
    visualize_pointcloud(points, labels)


if __name__ == '__main__':
    main()
